# rag_kg.py

import os
import sys
import argparse
import logging

# 预先设置NLTK环境变量，避免重复下载
# rag_kg.py 中的 pre_setup_environment 函数

def pre_setup_environment():
    """预先设置环境变量"""
    # 设置NLTK数据目录
    nltk_data_dir = os.path.join(os.path.expanduser('~'), 'nltk_data')
    os.makedirs(nltk_data_dir, exist_ok=True)
    os.environ['NLTK_DATA'] = nltk_data_dir
    
    # 禁用ChromaDB遥测
    os.environ['ANONYMIZED_TELEMETRY'] = 'False'
    
    # 设置第三方库日志级别，减少不必要的输出
    # 第三方库设置为WARNING级别，这样文件会记录WARNING和ERROR，但不会记录DEBUG
    third_party_loggers = [
        'chromadb', 'urllib3', 'backoff', 'openai', 'dashscope', 
        'llama_index', 'httpx', 'asyncio', 'charset_normalizer'
    ]
    for logger_name in third_party_loggers:
        logging.getLogger(logger_name).setLevel(logging.WARNING)
    
    # 添加NLTK路径
    import nltk
    nltk.data.path.insert(0, nltk_data_dir)

# 执行预设置
pre_setup_environment()

# 添加 knowledge 到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "knowledge"))

from knowledge.core.kb_application import KnowledgeBaseApplication
from knowledge.database.db_manager import DatabaseManager
from knowledge.database.db_schema_manager import DBSchemaManager

def manage_database(rdb=False, rdata=False, rschema=False, rindex=False):
    """
    数据库管理功能
    
    Args:
        rdb: 是否重新构建数据库结构
        rdata: 是否重新加载数据
        rschema: 是否重新构建数据库描述文件
        rindex: 是否重新构建索引
    """
    print("🔄 开始数据库管理操作...")
    
    try:
        # 创建数据库管理器
        db_manager = DatabaseManager(verbose=True)
        
        # 创建Schema管理器
        schema_manager = DBSchemaManager(verbose=True)
        
        # 1. 数据库重构操作
        if rdb:
            print("🗑️  重构数据库（删除并重新创建）...")
            db_manager.reinit_DB()
            print("✅ 数据库重构完成")
        
        # 2. 重新加载数据操作
        if rdata:
            print("🔄 重新加载数据...")
            db_manager.reload_data()
            print("✅ 数据重新加载完成")
        
        # 3. 刷新数据库Schema
        if rschema:
            print("📄 刷新数据库Schema描述文件...")
            success = schema_manager.refresh_schema(db_manager)
            if success:
                print("✅ 数据库Schema刷新完成")
            else:
                print("❌ 数据库Schema刷新失败")
                return False
        
        # 4. 向量化存储Schema（构建索引）
        if rindex:
            print("🔍 将数据库Schema向量化存储到ChromaDB...")
            print("📝 使用DashScope TEXT_EMBEDDING_V2 嵌入模型...")
            
            # 使用DBSchemaManager的向量化方法
            success = schema_manager.vectorize_schema()
            
            if success:
                print("✅ 数据库Schema向量化存储完成")
            else:
                print("❌ 数据库Schema向量化存储失败")
                return False
        
        # 显示Schema信息
        print("\n🔍 Schema信息:")
        schema_info = schema_manager.get_schema_info()
        print(f"  📁 Schema文件: {schema_info.get('schema_file_path', 'N/A')}")
        print(f"  📊 文件存在: {'✅' if schema_info.get('schema_file_exists') else '❌'}")
        print(f"  🗂️  集合名称: {schema_info.get('collection_name', 'N/A')}")
        print(f"  📈 集合存在: {'✅' if schema_info.get('collection_exists') else '❌'}")
        print(f"  📄 文档数量: {schema_info.get('document_count', 0)}")
        print(f"  🤖 嵌入模型: {schema_info.get('embedding_model', 'N/A')}")
        
        # 显示分块统计信息
        if "chunk_statistics" in schema_info:
            stats = schema_info["chunk_statistics"]
            print(f"  📊 分块统计: 表块={stats.get('table_chunks', 0)}, 全局信息={stats.get('global_info_chunks', 0)}, 关系图={stats.get('relationship_chunks', 0)}")
        
        print("🎉 数据库管理操作完成!")
        return True
        
    except Exception as e:
        print(f"❌ 数据库管理操作失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def integrated_query(query_text, show_details=False):
    """融合查询：智能结合数据库和知识库查询"""
    print(f"🧠 融合查询: {query_text}")
    if show_details:
        print("📊 显示详细检索信息")
    
    # 创建应用实例
    kb_app = KnowledgeBaseApplication("enterprise_knowledge")
    
    # 初始化系统
    if not kb_app.initialize():
        print("❌ 系统初始化失败")
        return
    
    # 确保查询引擎已加载
    if not kb_app.ensure_query_engine_loaded():
        print("❌ 查询引擎加载失败")
        return
    
    # 执行融合查询
    try:
        result = kb_app.integrated_query(query_text, show_details=show_details)
        
        # 显示查询分析结果
        # if "analysis" in result:
        #     analysis = result["analysis"]
        #     print(f"🔍 问题分析:")
        #     print(f"  📊 数据库查询: {analysis.get('db_query', '无')}")
        #     print(f"  📚 知识库查询: {analysis.get('kb_query', '无')}")
        
        # # 显示查询过程
        # if "process_details" in result:
        #     process = result["process_details"]
            
        #     if process.get("database"):
        #         db_info = process["database"]
        #         print(f"🗃️  数据库查询结果:")
        #         print(f"   查询: {db_info.get('query', '无')}")
        #         print(f"   SQL: {db_info.get('sql_used', '无')}")
        #         print(f"   记录数: {db_info.get('result_count', 0)}")
        #         if db_info.get('results'):
        #             print(f"   示例数据: {db_info['results'][:2]}")  # 显示前2条记录
            
        #     if process.get("knowledge_base"):
        #         kb_info = process["knowledge_base"]
        #         print(f"📚 知识库查询结果:")
        #         print(f"   查询: {kb_info.get('query', '无')}")
        #         print(f"   来源数: {kb_info.get('source_count', 0)}")
        #         if kb_info.get('answer'):
        #             answer_preview = kb_info['answer'][:200] + "..." if len(kb_info['answer']) > 200 else kb_info['answer']
        #             print(f"   答案预览: {answer_preview}")
        
        # # 显示最终答案
        # print(f"\n💡 最终答案:")
        # print(f"   {result.get('answer', '无答案')}")
        
        return result
        
    except Exception as e:
        print(f"❌ 融合查询失败: {e}")
        import traceback
        traceback.print_exc()

def knowledge_base_query(query_text, show_details=False):
    """纯知识库查询"""
    print(f"📚 知识库查询: {query_text}")
    if show_details:
        print("📊 显示详细检索信息")
    
    # 创建应用实例
    kb_app = KnowledgeBaseApplication("enterprise_knowledge")
    
    # 初始化系统
    if not kb_app.initialize():
        print("❌ 系统初始化失败")
        return
    
    # 确保查询引擎已加载
    if not kb_app.ensure_query_engine_loaded():
        print("❌ 查询引擎加载失败")
        return
    
    # 执行纯知识库查询（通过smart_query禁用数据库查询）
    try:
        result = kb_app.smart_query(query_text, use_database=False, show_details=show_details)
        
        print(f"📚 查询来源: 知识库")
        
        if show_details and "metadata" in result:
            metadata = result["metadata"]
            if "retrieved_nodes" in metadata:
                print(f"📄 检索到 {metadata.get('node_count', 0)} 个相关文档")
        
        # 显示答案
        print(f"\n💡 答案:")
        print(f"   {result.get('answer', '无答案')}")
        
        return result
        
    except Exception as e:
        print(f"❌ 知识库查询失败: {e}")
        import traceback
        traceback.print_exc()

def database_query(query_text, show_details=False):
    """纯数据库查询"""
    print(f"🗃️  数据库查询: {query_text}")
    if show_details:
        print("📊 显示详细检索信息")
    
    # 创建应用实例
    kb_app = KnowledgeBaseApplication("enterprise_knowledge")
    
    # 初始化系统
    if not kb_app.initialize():
        print("❌ 系统初始化失败")
        return
    
    # 确保查询引擎已加载
    if not kb_app.ensure_query_engine_loaded():
        print("❌ 查询引擎加载失败")
        return
    
    # 执行纯数据库查询 - 直接调用数据库查询方法
    try:
        # 直接使用查询引擎的数据库查询方法，避免融合查询
        result = kb_app.query_engine.query_with_database(query_text, show_details=show_details)
        
        print(f"🗃️  查询来源: 数据库")
        
        if show_details and "metadata" in result:
            metadata = result["metadata"]
            print(f"📊 使用的SQL: {metadata.get('sql_query', 'N/A')}")
            print(f"📈 结果数量: {metadata.get('result_count', 0)}")
            if metadata.get('raw_results'):
                print(f"🔍 示例结果: {metadata['raw_results'][:2]}")
        
        # 显示完整答案，不截断
        print(f"\n💡 答案:")
        print(f"   {result.get('answer', '无答案')}")
        
        return result
        
    except Exception as e:
        print(f"❌ 数据库查询失败: {e}")
        import traceback
        traceback.print_exc()

def interactive_chat(use_integrated=True, show_details=False):
    """交互式聊天"""
    mode = "融合查询" if use_integrated else "知识库查询"
    print(f"💬 进入交互式聊天模式 - {mode} (输入 'quit' 退出)")
    print("💡 提示: 输入 'db:' 开头进行数据库查询，输入 'kg:' 开头进行知识库查询")
    if show_details:
        print("📊 显示详细检索信息")
    
    # 创建应用实例
    kb_app = KnowledgeBaseApplication("enterprise_knowledge")
    
    # 初始化系统
    if not kb_app.initialize():
        print("❌ 系统初始化失败")
        return
    
    # 检查索引是否存在
    collection_info = kb_app.get_collection_info()
    if not collection_info.get('exists', False):
        print("❌ 知识库索引不存在，请先运行构建命令")
        return
    
    while True:
        try:
            query = input("\n❓ 请输入问题: ").strip()
            if query.lower() in ['quit', 'exit', '退出']:
                break
            if not query:
                continue
            
            # 判断查询模式
            if query.startswith('db:'):
                # 数据库查询
                db_query = query[3:].strip()
                if db_query:
                    database_query(db_query, show_details)
                else:
                    print("❌ 请输入数据库查询内容")
                    
            elif query.startswith('kg:'):
                # 知识库查询
                kg_query = query[3:].strip()
                if kg_query:
                    knowledge_base_query(kg_query, show_details)
                else:
                    print("❌ 请输入知识库查询内容")
                    
            else:
                # 默认使用融合查询或指定模式
                if use_integrated:
                    integrated_query(query, show_details)
                else:
                    knowledge_base_query(query, show_details)
            
        except KeyboardInterrupt:
            print("\n👋 再见!")
            break
        except Exception as e:
            print(f"❌ 查询失败: {e}")

def build_knowledge_base(force_rebuild=False):
    """构建知识库"""
    print("🚀 开始构建企业知识库...")
    
    # 创建应用实例
    kb_app = KnowledgeBaseApplication("enterprise_knowledge")
    
    # 初始化系统
    print("📝 初始化系统...")
    if not kb_app.initialize():
        print("❌ 系统初始化失败")
        return False
    
    # 构建知识库
    knowledge_root = "./resource"
    print(f"📁 知识库路径: {os.path.abspath(knowledge_root)}")
    
    if not os.path.exists(knowledge_root):
        print("❌ 知识库目录不存在")
        return False
    
    print("🔨 构建知识库索引...")
    stats = kb_app.build_knowledge_base(knowledge_root, force_rebuild=force_rebuild)
    
    print(f"\n📊 构建统计:")
    print(f"   总文件数: {stats.get('total_files', 0)}")
    print(f"   总文档数: {stats.get('total_documents', 0)}")
    print(f"   文件夹: {', '.join(stats.get('folders', []))}")
    
    if stats.get('total_documents', 0) > 0:
        print("✅ 知识库构建成功!")
        return True
    else:
        print("❌ 知识库构建失败，未找到有效文档")
        return False

def setup_knowledge_logging(verbose=False):
    """设置知识库系统日志"""
    from knowledge import setup_logging
    setup_logging(verbose=verbose)
    
def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="企业知识库管理系统")
    parser.add_argument("action", nargs="?", choices=["build", "query", "db-query", "kg-query", "chat", "manage-db"], 
                       help="操作类型: build-构建知识库, query-融合查询, db-query-数据库查询, kg-query-知识库查询, chat-交互式聊天, manage-db-数据库管理")
    parser.add_argument("--query", "-q", help="查询内容")
    parser.add_argument("--force", "-f", action="store_true", help="强制重建索引")
    parser.add_argument("--verbose", "-v", action="store_true", help="显示详细检索信息")
    parser.add_argument("--kg-only", action="store_true", help="chat模式下仅使用知识库查询")
    
    # 简化的数据库管理参数
    parser.add_argument("--rdb", action="store_true", help="重新构建数据库结构")
    parser.add_argument("--rdata", action="store_true", help="重新加载数据")
    parser.add_argument("--rschema", action="store_true", help="重新构建数据库描述文件")
    parser.add_argument("--rindex", action="store_true", help="重新构建索引")
    
    args = parser.parse_args()
    # 确保调用知识库的日志设置
    setup_knowledge_logging(args.verbose)
    
    if args.action == "build":
        build_knowledge_base(force_rebuild=args.force)
    
    elif args.action == "query":
        if args.query:
            integrated_query(args.query, show_details=args.verbose)
        else:
            print("❌ 请使用 --query 参数指定查询内容")
    
    elif args.action == "db-query":
        if args.query:
            database_query(args.query, show_details=args.verbose)
        else:
            print("❌ 请使用 --query 参数指定查询内容")
    
    elif args.action == "kg-query":
        if args.query:
            knowledge_base_query(args.query, show_details=args.verbose)
        else:
            print("❌ 请使用 --query 参数指定查询内容")
    
    elif args.action == "chat":
        interactive_chat(use_integrated=not args.kg_only, show_details=args.verbose)
    
    elif args.action == "manage-db":
        manage_database(
            rdb=args.rdb,
            rdata=args.rdata,
            rschema=args.rschema,
            rindex=args.rindex
        )
    
    else:
        # 如果没有指定动作，显示帮助信息
        print("🤖 企业知识库管理系统")
        print("\n📚 知识库操作:")
        print("  python rag_kg.py build                    # 构建知识库")
        print("  python rag_kg.py build --force           # 强制重建知识库")
        print("  python rag_kg.py query -q \"问题\"         # 融合查询（智能结合数据库和知识库）")
        print("  python rag_kg.py query -q \"问题\" -v      # 融合查询（详细模式）")
        print("  python rag_kg.py db-query -q \"问题\"      # 纯数据库查询")
        print("  python rag_kg.py kg-query -q \"问题\"      # 纯知识库查询")
        print("  python rag_kg.py chat                    # 交互式聊天（默认融合查询）")
        print("  python rag_kg.py chat --kg-only          # 交互式聊天（仅知识库查询）")
        print("  python rag_kg.py chat -v                 # 交互式聊天（详细模式）")
        
        print("\n🗃️  数据库管理:")
        print("  python rag_kg.py manage-db --rdb --rdata --rschema --rindex")
        print("    # 完全重新构建数据库 (结构+数据+描述+索引)")
        print("  python rag_kg.py manage-db --rdata --rschema --rindex")
        print("    # 重新加载数据并更新描述和索引")
        print("  python rag_kg.py manage-db --rschema --rindex")
        print("    # 重新构建描述文件和索引")
        
        print("\n💡 使用示例:")
        print("  python rag_kg.py build")
        print('  python rag_kg.py query -q "黄睿负责的产品的详细介绍"')
        print('  python rag_kg.py query -q "商机金额统计和销售流程规范" -v')
        print('  python rag_kg.py db-query -q "2025年项目统计"')
        print('  python rag_kg.py kg-query -q "差旅报销政策"')
        print("  python rag_kg.py chat")
        print("  python rag_kg.py chat --kg-only")
        print("  python rag_kg.py manage-db --rindex")

if __name__ == "__main__":
    main()